! *data_fit* module
!
! Purpose:
!   Collect several algorithm for doing data fitting task
!
! Author:
!   Name: DONG Li
!   Email: dongli@lasg.iap.ac.cn

module data_fit

    use polynomial_basis
    use message
    use run_manager

    implicit none

    private

    public data_fit_init
    public data_fit_approximate
    public data_fit_evaluate
    public data_fit_derivative
    public data_fit_finalize
    public data_fit_min_num_data

    character(20) :: method
    integer :: num_dim

    logical :: initialized = .false.

    interface data_fit_evaluate
        module procedure data_fit_evaluate_single
        module procedure data_fit_evaluate_batch
    end interface data_fit_evaluate

contains

    subroutine data_fit_init(method_, num_dim_, max_exponent)
        character(*), intent(in) :: method_
        integer, intent(in) :: num_dim_, max_exponent

        if(initialized) return

        method = method_
        num_dim = num_dim_

        call polynomial_basis_init(num_dim, max_exponent)

        if(method == "findfit") then
            call mathlink_init
        end if

        initialized = .true.

        return
    end subroutine data_fit_init

    integer function data_fit_min_num_data(order) result(res)
        integer, intent(in) :: order

        res = polynomials(order)%num_monomial

        return
    end function data_fit_min_num_data

    subroutine data_fit_approximate(num_data, x, y, order, c)
        integer, intent(in) :: num_data, order
        real, intent(in) :: x(num_dim,num_data), y(num_data)
        real, intent(out), allocatable :: c(:)

        select case(method)
        case("dtrnlsp")
            call intel_dtrnlsp(num_data, x, y, order, c)
        case("findfit")
            call mathematica_findfit(num_data, x, y, order, c)
        case default
            call message_show(ERROR_MSG, &
                "data_fit",              &
                "data_fit_approximate",  &
                "invalid method """ // trim(method) // """")
            call run_manager_end_run
        end select

        return
    end subroutine data_fit_approximate

    real function data_fit_evaluate_single(x, order, c) result(y)
        real, intent(in) :: x(num_dim)
        integer, intent(in) :: order
        real, intent(in) :: c(*)

        y = polynomial_basis_evaluate(x, order, c)

        return
    end function data_fit_evaluate_single

    subroutine data_fit_evaluate_batch(num_data, x, y, order, c)
        integer, intent(in) :: num_data, order
        real, intent(in) :: x(num_dim,num_data)
        real, intent(out) :: y(num_data)
        real, intent(in) :: c(*)

        integer :: i

        do i = 1, num_data
            y(i) = data_fit_evaluate_single(x(:,i), order, c)
        end do

        return
    end subroutine data_fit_evaluate_batch

    real function data_fit_derivative(x, order, c, l, o) result(d)
        real, intent(in) :: x(num_dim)
        integer, intent(in) :: order
        real, intent(in) :: c(*)
        integer, intent(in) :: l ! dimension of differentiation
        integer, intent(in) :: o ! order of differentiation

        d = polynomial_basis_derivative(x, order, c, l, o)

        return
    end function data_fit_derivative

    subroutine intel_dtrnlsp(num_data, x, y, order, c)
        include "mkl_rci.fi"

        integer, intent(in) :: num_data, order
        real, intent(in) :: x(num_dim,num_data), y(num_data)
        real, intent(out), allocatable :: c(:)

        ! internal variables for intel dtrnlsp_* routines
        double precision :: eps(6)
        double precision :: jac_eps
        integer :: rci_request
        double precision :: F(num_data)
        double precision :: J(num_data,polynomials(order)%num_monomial)
        integer :: iter
        integer :: st_cr
        logical :: flag
        integer :: max_iter1, max_iter2
        double precision :: rs
        double precision :: r1, r2
        integer(8) :: handle
        integer :: i, res

        if(allocated(c) == .false.) then
            allocate(c(polynomials(order)%num_monomial))
            c = 0.0
        end if

        eps(1) = 1.d-5
        eps(2) = 1.d-5
        eps(3) = 1.d-5
        eps(4) = 1.d-5
        eps(5) = 1.d-5
        eps(6) = 1.d-5

        max_iter1 = 1000
        max_iter2 = 100
        rs = 100.d0
        jac_eps = 1.d-8

        res = dtrnlsp_init(handle, polynomials(order)%num_monomial, &
            num_data, c, eps, max_iter1, max_iter2, rs)
        if(res == tr_invalid_option) then
            call message_show(ERROR_MSG, &
                "data_fit",              &
                "intel_dtrnlsp",         &
                "dtrnlsp_init says invalid option, check Intel MKL")
            stop
        else if(res == tr_out_of_memory) then
            call message_show(ERROR_MSG, &
                "data_fit",              &
                "intel_dtrnlsp",         &
                "dtrnlsp_init says out of memory, check Intel MKL")
            stop
        end if
        if(res /= tr_success) then
            call mkl_freebuffers
            stop
        end if

        rci_request = 0
        flag = .true.
        do while(flag)
            res = dtrnlsp_solve(handle, F, J, rci_request)
            if(res /= tr_success) then
                call message_show(ERROR_MSG, &
                    "data_fit",              &
                    "intel_dtrnlsp",         &
                    "dtrnlsp_solve encounters internal error, check Intel MKL")
                call mkl_freebuffers
                stop
            end if
            select case(rci_request)
            case(-1, -2, -3, -4, -5, -6)
                flag = .false.
            case(1)
                call obj_func(num_data, polynomials(order)%num_monomial, &
                    dble(c), F)
            case(2)
                res = djacobi(obj_func, polynomials(order)%num_monomial, &
                    num_data, J, dble(c), jac_eps)
                if(res /= tr_success) then
                    call message_show(ERROR_MSG, &
                        "data_fit",              &
                        "intel_dtrnlsp",         &
                        "djacobi encounters internal error, check Intel MKL")
                    call mkl_freebuffers
                    stop
                end if
            end select
        end do

        call mkl_freebuffers
        
        return

        contains

            subroutine obj_func(num_data, num_monomial, c, F)
                integer, intent(in) :: num_data, num_monomial
                double precision, intent(in) :: c(num_monomial)
                double precision, intent(out) :: F(num_data)

                integer :: i, j, k
                double precision :: temp

                do k = 1, num_data
                    F(k) = y(k)
                    do j = 1, num_monomial
                        temp = 1.d0
                        do i = 1, num_dim
                            temp = temp*x(i,k)**polynomials(order)%exponents(j,i)
                        end do
                        F(k) = F(k)-c(j)*temp
                    end do
                end do

                return
            end subroutine obj_func

    end subroutine intel_dtrnlsp

    subroutine mathematica_findfit(num_data, x, f, order, c)
        integer, intent(in) :: num_data, order
        real, intent(in) :: x(num_dim,num_data), f(num_data)
        real, intent(out), allocatable :: c(:)

        integer :: res

        if(allocated(c)) deallocate(c)
        allocate(c(polynomials(order)%num_monomial))

        ! call C program for linking Mathematica
        call findfit(num_dim, polynomials(order)%num_monomial, &
            num_data, x, f, c, polynomials(order)%exponents)

        return
    end subroutine mathematica_findfit

    subroutine data_fit_finalize

        if(method == "findfit") then
            call mathlink_finalize
        end if

        return
    end subroutine data_fit_finalize

end module data_fit
